import torch
from torch.nn import Module
from torch.nn import Linear, Softmax
from torch import einsum
class InterestMatch(Module):
    def __init__(self, dim, out_dim):
        super().__init__()
        self.layer1 = Linear(dim, out_dim, bias=False)
        self.softmax = Softmax(dim=-1)

    def forward(self, input, input_guide):
        input = self.layer1(input)
        if len(input.shape) == 4:
            s = einsum('bsfd,bd->bsf', input, input_guide)
        elif len(input.shape) == 3:
            s = einsum('bsd,bd->bs', input, input_guide)

        similary = self.softmax(s)
        out = similary.unsqueeze(-1) * input

        return out.sum(-2)